from ultralytics import YOLO #在代码中导入yolo模块

#导入训练好的模型best.pt
model=YOLO('best.pt')

#随意找一些测试数据
#图片数据和视频数据都可以，直接将数据传入接口就可以了
model.predict("xxx.jpg",save=True)
model.predict("xxx.mp4",save=True)

#自己构造一些数据
#在识别自己构造的数据时，传入了classes=[0,2],
#代表只输出0和2，也就是安全帽是否佩戴这两类
#line_width=30表示指定识别框的字体大小为30
model.predict('xxx1.jpg',save=True,classes=[0,2],line_width=30)
model.predict("xxx2.jpg",save=True,classes=[0,2],line_width=30)
model.predict("xxx1.MP4",save=True,classes=[0,2],line_width=30)
model.predict("xxx2.MP4",save=True,classes=[0,2],line_width=30)